from vanna.openai import OpenAI_Chat
from vanna.chromadb import ChromaDB_VectorStore
from openai import OpenAI

class MyVanna(ChromaDB_VectorStore, OpenAI_Chat):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        
        # 创建自定义的 OpenAI 客户端指向 DeepSeek
        client = OpenAI(
            api_key='your_api_key',	#注意替换成自己在MaaS平台注册的API Key。
            base_url='https://api.modelarts-maas.com/v1',  #注意替换成自己在MaaS平台获取到的的API地址。
        )
        
        OpenAI_Chat.__init__(self, client=client, config=config)

# 初始化
vn = MyVanna(config={
    'model': 'DeepSeek-V3'
})

vn.connect_to_mysql(host='localhost', dbname='school', user='root', password='your_password', port=3306)

vn.train(ddl="""
CREATE TABLE school.students (
    student_id INT AUTO_INCREMENT PRIMARY KEY COMMENT '学生ID',
    name VARCHAR(50) NOT NULL COMMENT '学生姓名',
    age INT COMMENT '年龄',
    major VARCHAR(100) COMMENT '专业',
    email VARCHAR(100) UNIQUE COMMENT '电子邮件'
);
CREATE TABLE school.courses (
    course_id INT AUTO_INCREMENT PRIMARY KEY COMMENT '课程ID',
    course_name	 VARCHAR(100) NOT NULL COMMENT '课程名称',
    credits INT COMMENT '学分',
    teacher_id INT COMMENT '教师ID'
);
CREATE TABLE school.enrollments (
    enrollment_id INT AUTO_INCREMENT PRIMARY KEY COMMENT '选课记录ID',
    student_id INT NOT NULL COMMENT '学生ID',
    course_id INT NOT NULL COMMENT '课程ID',
    grade FLOAT COMMENT '成绩'
);
""")

from vanna.flask import VannaFlaskApp

VannaFlaskApp(vn, allow_llm_to_see_data=True).run()